from rest_framework.filters import BaseFilterBackend


class MenuFilter(BaseFilterBackend):
    def filter_queryset(self, request, queryset, view):
        user = request.user
        if view.action == 'list':
            if user.is_superuser:
                # 如果是超级管理员 返回所有主菜单
                return queryset.filter(type=0)
            # 获取当前用户的role的menu
            role_id_list = user.role.all().values_list('menu', flat=True)
            filter_queryset = queryset.filter(type=0, id__in=role_id_list)
            return filter_queryset
        elif view.action == 'menu_list':
            return queryset.filter(type=0)
        else:
            return queryset
